import numpy as np
from scipy.linalg import solve


def active_set_qp(Q, c, A, b, x0, max_iter=100, tol=1e-8):
    n = Q.shape[0]
    m = A.shape[0]

    x = x0.copy()
    active = [i for i in range(m) if abs(A[i] @ x - b[i]) < tol]
    W = active
    history = []
    lambd = np.zeros(len(W))

    for k in range(max_iter):
        # 1. 识别当前有效集
        # print(W)
        alpha = 1
        # 2. 解KKT系统：在有效集上求解方向d
        if W:
            AW = A[W]
            KKT_matrix = np.block([[Q, -AW.T], [AW, np.zeros((len(W), len(W)))]])
            rhs = -np.hstack((Q @ x + c, np.zeros(len(W))))
        else:
            KKT_matrix = Q
            rhs = -(Q @ x + c)

        try:
            sol = solve(KKT_matrix, rhs, assume_a="sym")
        except np.linalg.LinAlgError:
            print("线性系统不可解，终止")
            break

        if W:
            d = sol[:n]
            lambd = sol[n:]
        else:
            d = sol
            lambd = np.array(0)
        # print("W = ", W, d, lambd)
        print(x.copy(), W, d, lambd, alpha)
        history.append((x.copy(), W.copy(), d, lambd, alpha))
        # 3. 检查是否是极小点
        if np.linalg.norm(d) < tol:
            if W:
                # 看拉格朗日乘子是否都>=0
                if all(lambd >= -tol):
                    print("最优解找到！")
                    break
                else:
                    # 移除使拉格朗日乘子为负的约束
                    idx = np.argmin(lambd)
                    # print(idx)
                    W.pop(idx)
                    print("移除约束", idx)
            else:
                print("最优解找到！（无约束）")
                break
        else:
            # 4. 沿着方向d，找最大可行步长
            # alpha = 1.0
            blocking_constraint = None
            for i in range(m):
                if i not in W:
                    Ai_d = A[i] @ d
                    if Ai_d < -tol:
                        step = (b[i] - A[i] @ x) / Ai_d
                        if step < alpha:
                            alpha = step
                            blocking_constraint = i

            # 5. 更新x
            x = x + alpha * d

            # 6. 如果有新的阻挡约束，加入有效集
            if blocking_constraint is not None:
                W.append(blocking_constraint)

    return x, history


if __name__ == "__main__":
    Q = np.array([[4.0, 1.0, 0.0], [1.0, 3.0, 0], [0, 0, 2]])
    c = np.array([-8.0, -6, -4.0])
    A = np.array(
        [
            [-1, -2, -1],
            [-2, 0, -1],
            [0, -2, -1],
            [1, 0, 0],
            [0, 1, 0],
            [0, 0, 1],
        ]
    )
    b = np.array([-10, -8, -7, 0, 0, 0])
    x0 = np.array([0, 0, 0])
    p = None

    x_opt, history = active_set_qp(Q, c, A, b, x0)
    for i in range(len(history)):
        print(f"Iteration {i}: {history[i]}")
        p = input("")
    print("Optimal solution:", x_opt)
